#!/bin/bash

# 智能 APT 软件源更换脚本（修复版）
# 支持：Ubuntu, Debian, ARM64, x86_64
# 源选项：清华源、中科大源、阿里源

set -e

# 颜色定义
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
BLUE='\033[0;34m'
NC='\033[0m' # No Color

# 日志函数
log_info() {
    echo -e "${BLUE}[INFO]${NC} $1"
}

log_success() {
    echo -e "${GREEN}[SUCCESS]${NC} $1"
}

log_warning() {
    echo -e "${YELLOW}[WARNING]${NC} $1"
}

log_error() {
    echo -e "${RED}[ERROR]${NC} $1"
}

# 检测系统架构
detect_architecture() {
    local arch=$(uname -m)
    case $arch in
        x86_64|amd64)
            echo "amd64"
            ;;
        aarch64|arm64)
            echo "arm64"
            ;;
        armv7l|armhf)
            echo "armhf"
            ;;
        *)
            echo "unknown"
            ;;
    esac
}

# 获取 Ubuntu 代号
get_ubuntu_codename() {
    local version=$1
    case $version in
        "22.04") echo "jammy" ;;
        "20.04") echo "focal" ;;
        "18.04") echo "bionic" ;;
        "16.04") echo "xenial" ;;
        "24.04") echo "noble" ;;
        *) echo "unknown" ;;
    esac
}

# 获取 Debian 代号
get_debian_codename() {
    local version=$1
    case $version in
        "12") echo "bookworm" ;;
        "11") echo "bullseye" ;;
        "10") echo "buster" ;;
        "9") echo "stretch" ;;
        "8") echo "jessie" ;;
        "testing") echo "testing" ;;
        "sid") echo "sid" ;;
        *) echo "unknown" ;;
    esac
}

# 检测系统版本和代号
detect_distro() {
    if [ -f /etc/os-release ]; then
        . /etc/os-release
        local codename=""
        if [ "$ID" = "ubuntu" ]; then
            codename=$(get_ubuntu_codename "$VERSION_ID")
        elif [ "$ID" = "debian" ]; then
            if [ -n "$VERSION_CODENAME" ]; then
                codename="$VERSION_CODENAME"
            else
                codename=$(get_debian_codename "$VERSION_ID")
            fi
        fi
        echo "$ID $VERSION_ID $codename"
    else
        log_error "无法检测系统版本"
        exit 1
    fi
}

# 备份原有源列表
backup_sources() {
    local backup_dir="/etc/apt/backup_$(date +%Y%m%d_%H%M%S)"
    mkdir -p "$backup_dir"
    
    if [ -f /etc/apt/sources.list ]; then
        cp /etc/apt/sources.list "$backup_dir/"
        log_info "已备份 sources.list 到 $backup_dir"
    fi
    
    if [ -d /etc/apt/sources.list.d ]; then
        cp -r /etc/apt/sources.list.d "$backup_dir/" 2>/dev/null || true
        log_info "已备份 sources.list.d 到 $backup_dir"
    fi
}

# 生成Ubuntu源内容
generate_ubuntu_sources() {
    local codename=$1
    local arch=$2
    local source_url=$3
    local source_type=$4
    
    cat << EOF
# Ubuntu 软件源 - $source_type
deb ${source_url} ${codename} main restricted universe multiverse
# deb-src ${source_url} ${codename} main restricted universe multiverse

deb ${source_url} ${codename}-security main restricted universe multiverse
# deb-src ${source_url} ${codename}-security main restricted universe multiverse

deb ${source_url} ${codename}-updates main restricted universe multiverse
# deb-src ${source_url} ${codename}-updates main restricted universe multiverse

deb ${source_url} ${codename}-backports main restricted universe multiverse
# deb-src ${source_url} ${codename}-backports main restricted universe multiverse

# 预发布软件源，不建议启用
# deb ${source_url} ${codename}-proposed main restricted universe multiverse
# deb-src ${source_url} ${codename}-proposed main restricted universe multiverse
EOF
}

# 生成Debian源内容
generate_debian_sources() {
    local codename=$1
    local arch=$2
    local source_url=$3
    
    local component="main contrib non-free"
    if [ "$codename" = "sid" ] || [ "$codename" = "testing" ]; then
        component="main"
    fi
    
    cat << EOF
# Debian 软件源
deb ${source_url} ${codename} ${component}
deb-src ${source_url} ${codename} ${component}

deb ${source_url} ${codename}-security ${component}
deb-src ${source_url} ${codename}-security ${component}

deb ${source_url} ${codename}-updates ${component}
deb-src ${source_url} ${codename}-updates ${component}
EOF
}

# 更换软件源
change_sources() {
    local distro=$1
    local version=$2
    local codename=$3
    local arch=$4
    local source_type=$5
    
    local source_url=""
    case $source_type in
        "tsinghua")
            case $distro in
                "ubuntu") source_url="https://mirrors.tuna.tsinghua.edu.cn/ubuntu/" ;;
                "debian") source_url="https://mirrors.tuna.tsinghua.edu.cn/debian/" ;;
            esac
            ;;
        "ustc")
            case $distro in
                "ubuntu") source_url="https://mirrors.ustc.edu.cn/ubuntu/" ;;
                "debian") source_url="https://mirrors.ustc.edu.cn/debian/" ;;
            esac
            ;;
        "aliyun")
            case $distro in
                "ubuntu") source_url="https://mirrors.aliyun.com/ubuntu/" ;;
                "debian") source_url="https://mirrors.aliyun.com/debian/" ;;
            esac
            ;;
    esac
    
    if [ -z "$source_url" ]; then
        log_error "不支持的发行版: $distro"
        return 1
    fi
    
    log_info "正在更换为 $source_type 源..."
    log_info "发行版: $distro $version ($codename)"
    log_info "架构: $arch"
    log_info "源地址: $source_url"
    
    # 备份原有源
    backup_sources
    
    # 生成新的源文件
    case $distro in
        "ubuntu")
            generate_ubuntu_sources "$codename" "$arch" "$source_url" "$source_type" > /etc/apt/sources.list
            ;;
        "debian")
            generate_debian_sources "$codename" "$arch" "$source_url" > /etc/apt/sources.list
            ;;
        *)
            log_error "不支持的发行版: $distro"
            return 1
            ;;
    esac
    
    log_success "软件源更换完成"
}

# 清理有问题的 PPA
clean_broken_ppa() {
    log_info "正在清理有问题的 PPA..."
    
    # 查找所有 .list 文件
    for list_file in /etc/apt/sources.list.d/*.list; do
        if [ -f "$list_file" ]; then
            # 检查文件中是否有 404 错误的仓库
            if grep -q "rabbitvcs" "$list_file" || grep -q "404" "$list_file" 2>/dev/null; then
                log_warning "禁用有问题的仓库: $(basename $list_file)"
                mv "$list_file" "$list_file.disabled" 2>/dev/null || true
            fi
        fi
    done
}

# 更新软件包缓存
update_apt_cache() {
    log_info "正在更新软件包缓存..."
    
    # 先清理可能的缓存问题
    apt clean
    rm -rf /var/lib/apt/lists/*
    
    # 更新缓存
    if apt update; then
        log_success "软件包缓存更新成功"
    else
        log_error "软件包缓存更新失败，尝试修复..."
        clean_broken_ppa
        apt update
    fi
}

# 显示系统信息
show_system_info() {
    log_info "=== 系统信息 ==="
    log_info "架构: $(detect_architecture)"
    local distro_info=$(detect_distro)
    local distro=$(echo $distro_info | awk '{print $1}')
    local version=$(echo $distro_info | awk '{print $2}')
    local codename=$(echo $distro_info | awk '{print $3}')
    log_info "发行版: $distro $version ($codename)"
    log_info "主机名: $(hostname)"
    log_info "内核版本: $(uname -r)"
}

# 恢复备份
restore_backup() {
    if ls /etc/apt/backup_* 1> /dev/null 2>&1; then
        local latest_backup=$(ls -td /etc/apt/backup_* | head -1)
        cp "$latest_backup/sources.list" /etc/apt/sources.list
        if [ -d "$latest_backup/sources.list.d" ]; then
            rm -rf /etc/apt/sources.list.d/*
            cp -r "$latest_backup/sources.list.d"/* /etc/apt/sources.list.d/ 2>/dev/null || true
        fi
        log_success "已恢复备份的源配置: $(basename $latest_backup)"
        return 0
    else
        log_error "未找到备份文件"
        return 1
    fi
}

# 主函数
main() {
    # 检查root权限
    if [ "$EUID" -ne 0 ]; then
        log_error "请使用 root 权限运行此脚本"
        log_info "使用方法: sudo $0"
        exit 1
    fi
    
    # 显示系统信息
    show_system_info
    
    # 检测系统和架构
    local distro_info=$(detect_distro)
    local distro=$(echo $distro_info | awk '{print $1}')
    local version=$(echo $distro_info | awk '{print $2}')
    local codename=$(echo $distro_info | awk '{print $3}')
    local arch=$(detect_architecture)
    
    if [ "$codename" = "unknown" ]; then
        log_error "不支持的 Ubuntu/Debian 版本: $version"
        exit 1
    fi
    
    # 选择源
    log_info "请选择软件源:"
    echo "1) 清华源 (Tsinghua)"
    echo "2) 中科大源 (USTC)"
    echo "3) 阿里源 (Aliyun)"
    echo "4) 恢复默认源"
    read -p "请输入选择 [1-4]: " choice
    
    case $choice in
        1)
            change_sources "$distro" "$version" "$codename" "$arch" "tsinghua"
            ;;
        2)
            change_sources "$distro" "$version" "$codename" "$arch" "ustc"
            ;;
        3)
            change_sources "$distro" "$version" "$codename" "$arch" "aliyun"
            ;;
        4)
            if restore_backup; then
                update_apt_cache
            fi
            exit 0
            ;;
        *)
            log_error "无效选择"
            exit 1
            ;;
    esac
    
    # 清理有问题的 PPA
    clean_broken_ppa
    
    # 更新缓存
    update_apt_cache
    
    log_success "软件源更换完成！"
}

# 帮助信息
show_help() {
    echo "用法: $0 [选项]"
    echo "选项:"
    echo "  -h, --help    显示此帮助信息"
    echo "  -a, --auto    自动选择清华源"
    echo "  -u, --ustc    选择中科大源"
    echo "  -l, --aliyun  选择阿里源"
    echo ""
    echo "示例:"
    echo "  sudo $0          # 交互式选择"
    echo "  sudo $0 --auto   # 自动选择清华源"
}

# 命令行参数处理
case "${1:-}" in
    -h|--help)
        show_help
        exit 0
        ;;
    -a|--auto)
        # 自动模式使用清华源
        distro_info=$(detect_distro)
        distro=$(echo $distro_info | awk '{print $1}')
        version=$(echo $distro_info | awk '{print $2}')
        codename=$(echo $distro_info | awk '{print $3}')
        arch=$(detect_architecture)
        change_sources "$distro" "$version" "$codename" "$arch" "tsinghua"
        clean_broken_ppa
        update_apt_cache
        ;;
    -u|--ustc)
        distro_info=$(detect_distro)
        distro=$(echo $distro_info | awk '{print $1}')
        version=$(echo $distro_info | awk '{print $2}')
        codename=$(echo $distro_info | awk '{print $3}')
        arch=$(detect_architecture)
        change_sources "$distro" "$version" "$codename" "$arch" "ustc"
        clean_broken_ppa
        update_apt_cache
        ;;
    -l|--aliyun)
        distro_info=$(detect_distro)
        distro=$(echo $distro_info | awk '{print $1}')
        version=$(echo $distro_info | awk '{print $2}')
        codename=$(echo $distro_info | awk '{print $3}')
        arch=$(detect_architecture)
        change_sources "$distro" "$version" "$codename" "$arch" "aliyun"
        clean_broken_ppa
        update_apt_cache
        ;;
    *)
        main
        ;;
esac